from qkdKeyGen import QKDKeyGen
from qkdBuffer import QKDBuffer
from qkdKeyManageSystem import QKDKeyManagerSystem
import threading
import time

def start_kms_server(kms):
    # Run the server in a separate thread to simulate the environment
    server_thread = threading.Thread(target=kms.run_server)
    server_thread.daemon = True
    server_thread.start()

def test_key_generation():
    # Step 1: Create KMS instances with QKDBuffer
    kms1 = QKDKeyManagerSystem(kms_id="KMS001", max_keys_per_request=5, min_key_size=128, max_key_size=256, kms_hostname="kms001.example.com")
    kms2 = QKDKeyManagerSystem(kms_id="KMS002", max_keys_per_request=5, min_key_size=128, max_key_size=256, kms_hostname="kms002.example.com")

    # the max_capacity value and min_capacity shoule be the same in the same qkdKeyGen module
    buffer1 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="192.168.1.1")
    buffer2 = QKDBuffer(key_size=256, max_capacity=200, min_capacity=20, ip_address="192.168.1.2")

    kms1.add_buffer(buffer1)
    kms2.add_buffer(buffer2)

    # Step 2: Start the KMS servers
    start_kms_server(kms1)
    start_kms_server(kms2)
    
    # time.sleep(2)
    
    # Step 3: Generate QKD Keys
    key_gen = QKDKeyGen(key_size=256, key_rate=5, packet_size=1024, data_rate=100, max_sifting_packets=10, master_address="192.168.1.1", slave_address="192.168.1.2")
    key_gen.start_communication()

if __name__ == "__main__":
    test_key_generation()
